import cv2 as cv
import numpy as np
import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk


class DocScanner:
    def __init__(self):
        pass

    def load_image(self, file_path):
        """加载图片并找到文档的四个角落。"""
        img = cv.imread(file_path)
        if img is None:
            raise ValueError(f"无法加载图片: {file_path}")
            
        img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        thresh, binary_img = cv.threshold(img_gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
        
        # 找到二值图的所有轮廓
        contours, hierarchy = cv.findContours(binary_img, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        
        # 找到最大的周长的四边形的四个顶点
        max_perimeter = 0
        extreme_pnts = None
        for contour in contours:
            perimeter = cv.arcLength(contour, True)
            approx = cv.approxPolyDP(contour, 0.02 * perimeter, True)
            if len(approx) == 4 and perimeter > max_perimeter:
                max_perimeter = perimeter
                extreme_pnts = approx
        
        if extreme_pnts is None:
            raise ValueError("无法找到合适的四边形轮廓")
            
        corners = self.order_points(extreme_pnts.reshape(4, 2))
        return img, corners

    def order_points(self, pts):
        """对给定的四个点进行排序，返回排序后的点。"""
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]  # 左上角
        rect[2] = pts[np.argmax(s)]  # 右下角
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]  # 右上角
        rect[3] = pts[np.argmax(diff)]  # 左下角
        
        print(f"排序后的角点:")
        print(f"0: 左上角: {rect[0]}")
        print(f"1: 右上角: {rect[1]}")
        print(f"2: 右下角: {rect[2]}")
        print(f"3: 左下角: {rect[3]}")
        
        return rect

    def crop_image(self, img, corners):
        """根据四个角落裁剪图片。"""
        # 确保角点是正确的格式
        if not isinstance(corners, np.ndarray):
            corners = np.array(corners)
            
        # 确保角点是 float32 类型
        corners = corners.astype(np.float32)
        
        # 检查角点数量
        if corners.shape != (4, 2):
            raise ValueError(f"角点数量不正确，期望 (4, 2)，实际 {corners.shape}")
            
        # 提取四个角点
        top_left_corner = corners[0]
        top_right_corner = corners[1]
        bottom_right_corner = corners[2]
        bottom_left_corner = corners[3]

        # 计算目标图片的宽度和高度
        width, height = self.get_image_dimensions((top_left_corner, top_right_corner, bottom_right_corner, bottom_left_corner))  

        # 创建目标点的坐标
        dst_points = np.array([
            [0, 0],
            [width - 1, 0],
            [width - 1, height - 1],
            [0, height - 1]
        ], dtype=np.float32)
            
        # 通过投影变换展平文档
        M = cv.getPerspectiveTransform(corners, dst_points)
        dst = cv.warpPerspective(img, M, (width, height))
        return dst

    def get_image_dimensions(self, corners):
        """计算图片的宽度和高度。"""
        (tl, tr, br, bl) = corners
        
        # 计算新图像的宽度
        widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
        widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
        maxWidth = max(int(widthA), int(widthB))
        
        # 计算新图像的高度
        heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
        heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
        maxHeight = max(int(heightA), int(heightB))
        
        return maxWidth, maxHeight


class DocScannerApp:
    def __init__(self):
        self.doc_scanner = DocScanner()
        self.corners = None
        self.dragging_idx = -1
        self.img = None
        
        self.root = tk.Tk()
        self.root.title("文档扫描应用")
        
        self.canvas = tk.Canvas(self.root, width=600, height=400)
        self.canvas.pack()
        
        self.canvas.bind("<Button-1>", self.mouse_callback)
        self.canvas.bind("<ButtonRelease-1>", self.mouse_callback)
        self.canvas.bind("<Motion>", self.mouse_callback)

        btn_select = tk.Button(self.root, text="选择图片", command=self.select_image)
        btn_select.pack(side=tk.LEFT, padx=10, pady=10)

        self.show_mouse_move = tk.BooleanVar(value=False)
        chk_show_mouse_move = tk.Checkbutton(self.root, text="显示鼠标移动", variable=self.show_mouse_move)
        chk_show_mouse_move.pack(side=tk.LEFT, padx=10, pady=10)

        btn_crop = tk.Button(self.root, text="裁剪", command=self.crop)
        btn_crop.pack(side=tk.RIGHT, padx=10, pady=10)

    def select_image(self):
        """选择图片并加载图片和角落。"""
        file_path = filedialog.askopenfilename()
        if not file_path:
            return  # 用户取消选择
        
        try:
            self.img, self.corners = self.doc_scanner.load_image(file_path)
            print(f"成功加载图片，角点数量: {len(self.corners)}")
            
            img_height, img_width, _ = self.img.shape
            if img_height > 600 or img_width > 800:
                scale = min(600/img_height, 800/img_width)
                self.img = cv.resize(self.img, (int(img_width*scale), int(img_height*scale)))
                self.corners = (self.corners * scale).astype(int)
                img_height, img_width, _ = self.img.shape
                
            self.canvas.config(width=img_width, height=img_height)
            self.canvas.pack()
            self.root.geometry(f"{img_width+50}x{img_height+100}")
            self.redraw()
            
        except Exception as e:
            print(f"加载图片时出错: {str(e)}")
            tk.messagebox.showerror("错误", f"加载图片时出错: {str(e)}")

    def mouse_callback(self, event):
        """鼠标回调函数，处理鼠标按下、释放和移动事件。"""
        x, y = event.x, event.y
        if event.type == tk.EventType.ButtonPress:
            print(f"鼠标在({x},{y})按下")
            for idx, corner in enumerate(self.corners):
                if np.sqrt(((corner[0] - x)**2) + ((corner[1] - y)**2)) < 10:
                    self.dragging_idx = idx
                    print(f"开始拖动角落: corners[{idx}]={corner}")
                    break
        elif event.type == tk.EventType.ButtonRelease:
            print(f"鼠标在({x},{y})释放")
            self.dragging_idx = -1
        elif event.type == tk.EventType.Motion:
            if self.show_mouse_move.get():
                print(f"鼠标在({x},{y})移动")
            if self.dragging_idx != -1:
                self.corners[self.dragging_idx] = np.array([x, y])
                self.redraw()

    def redraw(self):
        """重新绘制图片和角落。"""
        if self.img is None or self.corners is None:
            return
            
        img_copy = self.img.copy()
        for idx, corner in enumerate(self.corners):
            cv.circle(img_copy, tuple(corner.astype(int)), 5, (0, 255, 0), -1)
            cv.putText(img_copy, str(idx), tuple(corner.astype(int) + 10), 
                      cv.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
        cv.polylines(img_copy, [self.corners], True, (0, 255, 0), 2)
        
        img_tk = self.cv2image_to_tkinter_image(img_copy)
        self.canvas.delete("all")  # 清除画布
        self.canvas.create_image(0, 0, anchor=tk.NW, image=img_tk)
        self.canvas.image = img_tk

    def cv2image_to_tkinter_image(self, cv2_image):
        """将 OpenCV 图片转换为 Tkinter 图片。"""
        cv2_image_rgb = cv.cvtColor(cv2_image, cv.COLOR_BGR2RGB)
        pil_image = Image.fromarray(cv2_image_rgb)
        return ImageTk.PhotoImage(pil_image)

    def crop(self):
        """裁剪图片。"""
        if self.img is None or self.corners is None:
            tk.messagebox.showerror("错误", "请先选择图片")
            return
            
        try:
            # 确保角点是 float32 类型
            corners_float32 = self.corners.astype(np.float32)
            
            cropped_img = self.doc_scanner.crop_image(self.img, corners_float32)
            cropped_img_tk = self.cv2image_to_tkinter_image(cropped_img)
            
            cropped_image_window = tk.Toplevel(self.root)
            cropped_image_window.title("裁剪后的图片")
            
            img_label = tk.Label(cropped_image_window, image=cropped_img_tk)
            img_label.pack(padx=10, pady=10)
            
            btn_save = tk.Button(cropped_image_window, text="保存图片", 
                                command=lambda: self.save_image(cropped_img))
            btn_save.pack(side=tk.LEFT, padx=10, pady=10)
            
            btn_close = tk.Button(cropped_image_window, text="关闭", 
                                command=cropped_image_window.destroy)
            btn_close.pack(side=tk.RIGHT, padx=10, pady=10)
            
            img_label.image = cropped_img_tk
            
        except Exception as e:
            print(f"裁剪图片时出错: {str(e)}")
            tk.messagebox.showerror("错误", f"裁剪图片时出错: {str(e)}")

    def save_image(self, img):
        """保存图片到文件"""
        file_path = filedialog.asksaveasfilename(
            defaultextension=".jpg",
            filetypes=[("JPEG files", "*.jpg"), ("PNG files", "*.png"), ("All files", "*.*")]
        )
        
        if file_path:
            try:
                cv.imwrite(file_path, img)
                tk.messagebox.showinfo("成功", f"图片已保存到: {file_path}")
            except Exception as e:
                tk.messagebox.showerror("错误", f"保存图片时出错: {str(e)}")


if __name__ == "__main__":
    app = DocScannerApp()
    app.run()